package ck

import (
	"context"
	"net"
	"strings"

	"github.com/go-sql-driver/mysql"
	"github.com/jackc/pgx/v4"
	"github.com/jackc/pgx/v4/stdlib"
	"github.com/jmoiron/sqlx"
	"golang.org/x/crypto/ssh"
)

// demo:
// 替换 username 为自己的名称
// 修改 dns 为数据库连接地址
// body, _ := ioutil.ReadFile(os.Getenv("HOME") + "/.ssh/id_rsa")
// client, _ := NewSSHClient("127.0.0.1:22", "username", string(body))
// db, _ := OpenDB("mysql", "dns", client)
// db.Select()

// OpenDB 打开数据库连接
func OpenDB(driver, source string, client ...*ssh.Client) (*sqlx.DB, error) {
	if len(client) > 0 {
		if driver == "mysql" {
			mysql.RegisterDialContext("tcpSSH", func(ctx context.Context, addr string) (conn net.Conn, err error) {
				return client[0].Dial("tcp", addr)
			})
			source = strings.Replace(source, "@tcp", "@tcpSSH", 1)
		} else if driver == "pgx" {
			c, err := pgx.ParseConfig(source)
			if err != nil {
				return nil, err
			}
			c.DialFunc = func(ctx context.Context, network, addr string) (conn net.Conn, err error) {
				return client[0].Dial(network, addr)
			}
			return sqlx.NewDb(stdlib.OpenDB(*c), driver), nil
		}
	}
	return sqlx.Open(driver, source)
}

// NewSSHClient 获取ssh客户端
func NewSSHClient(host, user, privateKey string) (*ssh.Client, error) {
	signer, err := ssh.ParsePrivateKey([]byte(privateKey))
	if err != nil {
		return nil, err
	}
	return ssh.Dial("tcp", host, &ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	})
}

// NewSSHPasswordClient 获取ssh客户端根据密码
func NewSSHPasswordClient(host, user, password string) (*ssh.Client, error) {
	return ssh.Dial("tcp", host, &ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{
			ssh.Password(password),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	})
}
